import argparse
import json
import time
import os
import gradio as gr
from pathlib import Path
from string import Template
import requests
from api.index \
   import login, list_motion_shop_tasks, list_common_materials, generate_motion_shop_video_upload_url, motion_shop_video_detect, \
          query_motion_shop_video_detect_result, submit_motion_shop_task, batch_query_motion_shop_task_status

# 用户登录
def gr_on_load(uuid):
  user_id = uuid
  if user_id:
    try:
        resp = login(user_id)
        if (resp['body']['Success'] == False):
           raise gr.exceptions.Error('用户登录失败，请刷新重试。')
        jwt = resp['body']['Data']['JwtToken']
        return [
          uuid_field.update(user_id),
          jwt_token_field.update(jwt)
        ]
    except (NameError, AttributeError):
        raise gr.exceptions.Error('用户登录失败，请刷新重试。')
  else:
      raise gr.exceptions.Error('用户登录失败，请刷新重试。')

# 页面js资源加载
def load_vite_js():
  with open('./vue-project/dist/assets/main.js', 'r') as file:
    vite_js = file.read().replace('\n', '')
    return Template('(u) => { ${vite_js}; return [u] }').substitute({'vite_js': vite_js})
  
# 加载内置形象数据
def load_figure_list():
   try:
      res = list_common_materials()
      if (res["Success"] == False):
         gr.Error("页面数据加载失败，请刷新页面重试")
         return
      return f'<script id="figure-list-script">{json.dumps(res["Data"])}</script>'
   except Exception as e:
      print('[load_figure_list]error:', str(e))
      gr.Error("页面数据加载失败，请刷新页面重试")
  
# 加载历史记录列表
def load_history_list(jwt, params):
   try:
      params = json.loads(params)
      res = list_motion_shop_tasks(jwt, params['current'], params['size'])
      if (res["Success"] == False):
         gr.Error('页面数据加载失败，请刷新页面重试')
      return f'<script id="history-list-result">{json.dumps(res)}</script>'
   except Exception as e:
      print('[load_history_list]error', str(e))
      gr.Error('页面数据加载失败，请刷新页面重试')
      return f'<script id="history-list-result">{json.dumps({"Success": False})}</script>'

# 提交项目
def create_project(jwt, params):
   print('[create_project]params:', params)
   params = json.loads(params)
   try:
      videoId = params["videoId"]
      avatarId = params["avatarId"]
      title = f'魔搭项目_{int(time.time())}'
      res = submit_motion_shop_task(jwt, title, videoId, avatarId)
      if (res["Success"] == False):
         gr.Error('任务提交失败，请重试')
      return f'<script id="create-result-field">{json.dumps(res)}</script>'
   except Exception as e:
      gr.Error('任务提交失败，请重试')
      print('[create_project]error:', str(e))
      return f'<script id="create-result-field">{json.dumps({ "Success": False})}</script>'

# 更新项目状态
def refresh_projects_status(jwt, taskList):
   try:
      taskList = json.loads(taskList)
      response = batch_query_motion_shop_task_status(jwt, taskIds=[task['TaskId'] for task in taskList])
      if (response["Success"] == True):
         return f'<script id="projects-status-update-result">{json.dumps(response["Data"])}</script>'
      return f'<script id="projects-status-update-result">{json.dumps({ "Tasks": [] })}</script>'
   except Exception as e:
      print('[refresh_projects_status]error:', str(e))
      return f'<script id="projects-status-update-result">{json.dumps({ "Tasks": [] })}</script>'

# 上传视频文件
def upload_file(jwt, file):
   if (file):
     with open(file.name, 'rb') as fileObj:
        try:
            upload_params = generate_motion_shop_video_upload_url(jwt)
            if (upload_params["Success"] != True):
               return f'<script id="upload-result-field">{json.dumps({"success": False, "error": "UploadFailed"})}</script>'
            upload_response = requests.put(
              upload_params["Data"]["UploadUrl"],
              data = fileObj.read()
            )
            if (upload_response.status_code < 200 or upload_response.status_code >= 300):
               return f'<script id="upload-result-field">{json.dumps({"success": False, "error": "UploadFailed"})}</script>'
            check_response = motion_shop_video_detect(jwt, upload_params["Data"]["OssKey"])
            if (check_response['Success'] != True):
               return f'<script id="upload-result-field">{json.dumps({"success": False, "error": "UploadFailed"})}</script>'
            return f'<script id="upload-result-field">{json.dumps({"success": True, "data": check_response["Data"]["JobId"]})}</script>'
        except Exception as e:
           print('[upload_file]error:', str(e))
           return f'<script id="upload-result-field">{json.dumps({"success": False, "error": "UploadFailed"})}</script>'
   return '<script id="upload-result-field"></script>'

# 开启视频检测      
def detect_video(jwt, jobId):
   while (True):
      try:
        time.sleep(1)
        check_response = query_motion_shop_video_detect_result(jwt, jobId)
        if (check_response["Success"] and (check_response["Data"]["Status"] == 'FINISHED' or check_response["Data"]["Status"] == 'FAILED')):
          return f'<script id="detect-result-field">{json.dumps(check_response["Data"])}</script>'
      except Exception as e:
         print('[detect_video]error:', str(e))
         continue
    
with gr.Blocks() as demo:
    # 用户登录信息
    uuid_field = gr.Text(label='modelscope_uuid', elem_id='uuid', visible=False)
    jwt_token_field = gr.Text(label='modelscope_jwt_token', elem_id='jwt_token', visible=False)

    # 页面内容
    input = gr.HTML(elem_id='app')
    
    # 上传按钮
    upload_result_field = gr.HTML('<script id="upload-result-field"></script>')
    upload_file_field = gr.File(elem_id='upload-file-field', visible=False, file_types=['video'])
    upload_file_field.change(
      fn=lambda : '<script id="upload-result-field"></script>', # 清空上次结果
      outputs=[upload_result_field],
      show_progress='hidden'
   ).then(
      fn=upload_file,
      inputs=[jwt_token_field, upload_file_field],
      outputs=[upload_result_field],
      show_progress='hidden'
    )

    # 视频检测按钮
    detect_result_field = gr.HTML('<script id="detect-result-field"></script>')
    detect_button = gr.Button(elem_id='gradio-detect-btn', visible=False)
    detect_button.click(
       fn=lambda : '<script id="detect-result-field"></script>', # 清空上次结果
       outputs=[detect_result_field],
       show_progress='hidden'
    ).then(
       fn=detect_video,
       inputs=[jwt_token_field, input],
       outputs=[detect_result_field],
       _js='(j, i) => [ j, window.getDetectJobId() ]',
       show_progress='hidden'
    )

    # 生成按钮
    create_result_field = gr.HTML(f'<script id="create-result-field"></script>')
    create_button = gr.Button(elem_id='gradio-create-btn', visible=False)
    create_button.click(
       fn=lambda : '<script id="create-result-field"></script>', # 清空上次结果
       outputs=[create_result_field],
       show_progress='hidden'
    ).then(
      fn=create_project,
      inputs=[jwt_token_field, input],
      outputs=[create_result_field],
      show_progress='hidden',
      _js='(j, i) => [j, window.getSubmitParams()]'
    )

    # 形象列表
    figure_list_field = gr.HTML('<script id="figure-list-script"></script>')
    figure_list_button = gr.Button(elem_id='gradio-figure-list-btn', visible=False)
    figure_list_button.click(
      fn=load_figure_list,
      inputs=None,
      outputs=[figure_list_field],
      show_progress='hidden'
    )

    # 历史记录
    history_list_field = gr.HTML('<script id="history-list-result"></script>')
    history_list_button = gr.Button(elem_id='gradio-history-list-btn', visible=False)
    history_list_button.click(
       fn=lambda : '<script id="history-list-result"></script>', # 清空上次结果
       outputs=[history_list_field],
       show_progress='hidden'
    ).then(
      fn=load_history_list,
      inputs=[jwt_token_field, input],
      outputs=[history_list_field],
      _js='(j) => [j, window.getHistoryListParams()]',
      show_progress='hidden'
    )

    # 历史项目状态更新
    projects_status_update_result_field = gr.HTML('<script id="projects-status-update-result"></script>')
    projects_status_update_button = gr.Button(elem_id='gradio-projects-status-update-btn', visible=False)
    projects_status_update_button.click(
       fn=lambda : '<script id="projects-status-update-result"></script>', # 清空上次结果
       outputs=[projects_status_update_result_field],
       show_progress='hidden'
    ).then(
      fn=refresh_projects_status,
      inputs=[jwt_token_field, input],
      outputs=[projects_status_update_result_field],
      _js='(j, i) => [j, window.getMakingProjectList()]',
      show_progress='hidden'
    )

    demo.load(fn=gr_on_load, inputs=[uuid_field], outputs=[uuid_field, jwt_token_field], _js=load_vite_js())

def parse_args():
   parser = argparse.ArgumentParser()
   parser.add_argument("--server_name", "-n", type=str, default="0.0.0.0", help="server name")
   parser.add_argument("--server_port", "-p", type=int, default=7861, help="server port")
   args = parser.parse_args()
   return args


def main():
   args = parse_args()
   server_name = args.server_name
   server_port = args.server_port
   demo.queue(concurrency_count=100).launch(server_name=server_name, server_port=server_port)


if __name__ == '__main__':
   main()
